// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Gael Guennebaud <gael.guennebaud@inria.fr>
// Copyright (C) 2007-2009 Benoit Jacob <jacob.benoit.1@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_DIAGONALMATRIX_H
#define EIGEN_DIAGONALMATRIX_H

namespace Eigen {

#ifndef EIGEN_PARSED_BY_DOXYGEN
template <typename Derived> class DiagonalBase : public EigenBase<Derived>
{
public:
    typedef typename internal::traits<Derived>::DiagonalVectorType DiagonalVectorType;
    typedef typename DiagonalVectorType::Scalar Scalar;
    typedef typename DiagonalVectorType::RealScalar RealScalar;
    typedef typename internal::traits<Derived>::StorageKind StorageKind;
    typedef typename internal::traits<Derived>::StorageIndex StorageIndex;

    enum
    {
        RowsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
        ColsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
        MaxRowsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
        MaxColsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
        IsVectorAtCompileTime = 0,
        Flags = NoPreferredStorageOrderBit
    };

    typedef Matrix<Scalar, RowsAtCompileTime, ColsAtCompileTime, 0, MaxRowsAtCompileTime, MaxColsAtCompileTime> DenseMatrixType;
    typedef DenseMatrixType DenseType;
    typedef DiagonalMatrix<Scalar, DiagonalVectorType::SizeAtCompileTime, DiagonalVectorType::MaxSizeAtCompileTime> PlainObject;

    EIGEN_DEVICE_FUNC
    inline const Derived& derived() const { return *static_cast<const Derived*>(this); }
    EIGEN_DEVICE_FUNC
    inline Derived& derived() { return *static_cast<Derived*>(this); }

    EIGEN_DEVICE_FUNC
    DenseMatrixType toDenseMatrix() const { return derived(); }

    EIGEN_DEVICE_FUNC
    inline const DiagonalVectorType& diagonal() const { return derived().diagonal(); }
    EIGEN_DEVICE_FUNC
    inline DiagonalVectorType& diagonal() { return derived().diagonal(); }

    EIGEN_DEVICE_FUNC
    inline Index rows() const { return diagonal().size(); }
    EIGEN_DEVICE_FUNC
    inline Index cols() const { return diagonal().size(); }

    template <typename MatrixDerived>
    EIGEN_DEVICE_FUNC const Product<Derived, MatrixDerived, LazyProduct> operator*(const MatrixBase<MatrixDerived>& matrix) const
    {
        return Product<Derived, MatrixDerived, LazyProduct>(derived(), matrix.derived());
    }

    typedef DiagonalWrapper<const CwiseUnaryOp<internal::scalar_inverse_op<Scalar>, const DiagonalVectorType>> InverseReturnType;
    EIGEN_DEVICE_FUNC
    inline const InverseReturnType inverse() const { return InverseReturnType(diagonal().cwiseInverse()); }

    EIGEN_DEVICE_FUNC
    inline const DiagonalWrapper<const EIGEN_EXPR_BINARYOP_SCALAR_RETURN_TYPE(DiagonalVectorType, Scalar, product)> operator*(const Scalar& scalar) const
    {
        return DiagonalWrapper<const EIGEN_EXPR_BINARYOP_SCALAR_RETURN_TYPE(DiagonalVectorType, Scalar, product)>(diagonal() * scalar);
    }
    EIGEN_DEVICE_FUNC
    friend inline const DiagonalWrapper<const EIGEN_SCALAR_BINARYOP_EXPR_RETURN_TYPE(Scalar, DiagonalVectorType, product)> operator*(const Scalar& scalar,
                                                                                                                                     const DiagonalBase& other)
    {
        return DiagonalWrapper<const EIGEN_SCALAR_BINARYOP_EXPR_RETURN_TYPE(Scalar, DiagonalVectorType, product)>(scalar * other.diagonal());
    }

    template <typename OtherDerived>
    EIGEN_DEVICE_FUNC
#ifdef EIGEN_PARSED_BY_DOXYGEN
        inline unspecified_expression_type
#else
        inline const DiagonalWrapper<const EIGEN_CWISE_BINARY_RETURN_TYPE(DiagonalVectorType, typename OtherDerived::DiagonalVectorType, sum)>
#endif
        operator+(const DiagonalBase<OtherDerived>& other) const
    {
        return (diagonal() + other.diagonal()).asDiagonal();
    }

    template <typename OtherDerived>
    EIGEN_DEVICE_FUNC
#ifdef EIGEN_PARSED_BY_DOXYGEN
        inline unspecified_expression_type
#else
        inline const DiagonalWrapper<const EIGEN_CWISE_BINARY_RETURN_TYPE(DiagonalVectorType, typename OtherDerived::DiagonalVectorType, difference)>
#endif
        operator-(const DiagonalBase<OtherDerived>& other) const
    {
        return (diagonal() - other.diagonal()).asDiagonal();
    }
};

#endif

/** \class DiagonalMatrix
  * \ingroup Core_Module
  *
  * \brief Represents a diagonal matrix with its storage
  *
  * \param _Scalar the type of coefficients
  * \param SizeAtCompileTime the dimension of the matrix, or Dynamic
  * \param MaxSizeAtCompileTime the dimension of the matrix, or Dynamic. This parameter is optional and defaults
  *        to SizeAtCompileTime. Most of the time, you do not need to specify it.
  *
  * \sa class DiagonalWrapper
  */

namespace internal {
    template <typename _Scalar, int SizeAtCompileTime, int MaxSizeAtCompileTime>
    struct traits<DiagonalMatrix<_Scalar, SizeAtCompileTime, MaxSizeAtCompileTime>>
        : traits<Matrix<_Scalar, SizeAtCompileTime, SizeAtCompileTime, 0, MaxSizeAtCompileTime, MaxSizeAtCompileTime>>
    {
        typedef Matrix<_Scalar, SizeAtCompileTime, 1, 0, MaxSizeAtCompileTime, 1> DiagonalVectorType;
        typedef DiagonalShape StorageKind;
        enum
        {
            Flags = LvalueBit | NoPreferredStorageOrderBit
        };
    };
}  // namespace internal
template <typename _Scalar, int SizeAtCompileTime, int MaxSizeAtCompileTime>
class DiagonalMatrix : public DiagonalBase<DiagonalMatrix<_Scalar, SizeAtCompileTime, MaxSizeAtCompileTime>>
{
public:
#ifndef EIGEN_PARSED_BY_DOXYGEN
    typedef typename internal::traits<DiagonalMatrix>::DiagonalVectorType DiagonalVectorType;
    typedef const DiagonalMatrix& Nested;
    typedef _Scalar Scalar;
    typedef typename internal::traits<DiagonalMatrix>::StorageKind StorageKind;
    typedef typename internal::traits<DiagonalMatrix>::StorageIndex StorageIndex;
#endif

protected:
    DiagonalVectorType m_diagonal;

public:
    /** const version of diagonal(). */
    EIGEN_DEVICE_FUNC
    inline const DiagonalVectorType& diagonal() const { return m_diagonal; }
    /** \returns a reference to the stored vector of diagonal coefficients. */
    EIGEN_DEVICE_FUNC
    inline DiagonalVectorType& diagonal() { return m_diagonal; }

    /** Default constructor without initialization */
    EIGEN_DEVICE_FUNC
    inline DiagonalMatrix() {}

    /** Constructs a diagonal matrix with given dimension  */
    EIGEN_DEVICE_FUNC
    explicit inline DiagonalMatrix(Index dim) : m_diagonal(dim) {}

    /** 2D constructor. */
    EIGEN_DEVICE_FUNC
    inline DiagonalMatrix(const Scalar& x, const Scalar& y) : m_diagonal(x, y) {}

    /** 3D constructor. */
    EIGEN_DEVICE_FUNC
    inline DiagonalMatrix(const Scalar& x, const Scalar& y, const Scalar& z) : m_diagonal(x, y, z) {}

#if EIGEN_HAS_CXX11
    /** \brief Construct a diagonal matrix with fixed size from an arbitrary number of coefficients. \cpp11
      * 
      * There exists C++98 anologue constructors for fixed-size diagonal matrices having 2 or 3 coefficients.
      * 
      * \warning To construct a diagonal matrix of fixed size, the number of values passed to this 
      * constructor must match the fixed dimension of \c *this.
      * 
      * \sa DiagonalMatrix(const Scalar&, const Scalar&)
      * \sa DiagonalMatrix(const Scalar&, const Scalar&, const Scalar&)
      */
    template <typename... ArgTypes>
    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DiagonalMatrix(const Scalar& a0, const Scalar& a1, const Scalar& a2, const ArgTypes&... args)
        : m_diagonal(a0, a1, a2, args...)
    {
    }

    /** \brief Constructs a DiagonalMatrix and initializes it by elements given by an initializer list of initializer
      * lists \cpp11
      */
    EIGEN_DEVICE_FUNC
    explicit EIGEN_STRONG_INLINE DiagonalMatrix(const std::initializer_list<std::initializer_list<Scalar>>& list) : m_diagonal(list) {}
#endif  // EIGEN_HAS_CXX11

    /** Copy constructor. */
    template <typename OtherDerived> EIGEN_DEVICE_FUNC inline DiagonalMatrix(const DiagonalBase<OtherDerived>& other) : m_diagonal(other.diagonal()) {}

#ifndef EIGEN_PARSED_BY_DOXYGEN
    /** copy constructor. prevent a default copy constructor from hiding the other templated constructor */
    inline DiagonalMatrix(const DiagonalMatrix& other) : m_diagonal(other.diagonal()) {}
#endif

    /** generic constructor from expression of the diagonal coefficients */
    template <typename OtherDerived> EIGEN_DEVICE_FUNC explicit inline DiagonalMatrix(const MatrixBase<OtherDerived>& other) : m_diagonal(other) {}

    /** Copy operator. */
    template <typename OtherDerived> EIGEN_DEVICE_FUNC DiagonalMatrix& operator=(const DiagonalBase<OtherDerived>& other)
    {
        m_diagonal = other.diagonal();
        return *this;
    }

#ifndef EIGEN_PARSED_BY_DOXYGEN
    /** This is a special case of the templated operator=. Its purpose is to
      * prevent a default operator= from hiding the templated operator=.
      */
    EIGEN_DEVICE_FUNC
    DiagonalMatrix& operator=(const DiagonalMatrix& other)
    {
        m_diagonal = other.diagonal();
        return *this;
    }
#endif

    /** Resizes to given size. */
    EIGEN_DEVICE_FUNC
    inline void resize(Index size) { m_diagonal.resize(size); }
    /** Sets all coefficients to zero. */
    EIGEN_DEVICE_FUNC
    inline void setZero() { m_diagonal.setZero(); }
    /** Resizes and sets all coefficients to zero. */
    EIGEN_DEVICE_FUNC
    inline void setZero(Index size) { m_diagonal.setZero(size); }
    /** Sets this matrix to be the identity matrix of the current size. */
    EIGEN_DEVICE_FUNC
    inline void setIdentity() { m_diagonal.setOnes(); }
    /** Sets this matrix to be the identity matrix of the given size. */
    EIGEN_DEVICE_FUNC
    inline void setIdentity(Index size) { m_diagonal.setOnes(size); }
};

/** \class DiagonalWrapper
  * \ingroup Core_Module
  *
  * \brief Expression of a diagonal matrix
  *
  * \param _DiagonalVectorType the type of the vector of diagonal coefficients
  *
  * This class is an expression of a diagonal matrix, but not storing its own vector of diagonal coefficients,
  * instead wrapping an existing vector expression. It is the return type of MatrixBase::asDiagonal()
  * and most of the time this is the only way that it is used.
  *
  * \sa class DiagonalMatrix, class DiagonalBase, MatrixBase::asDiagonal()
  */

namespace internal {
    template <typename _DiagonalVectorType> struct traits<DiagonalWrapper<_DiagonalVectorType>>
    {
        typedef _DiagonalVectorType DiagonalVectorType;
        typedef typename DiagonalVectorType::Scalar Scalar;
        typedef typename DiagonalVectorType::StorageIndex StorageIndex;
        typedef DiagonalShape StorageKind;
        typedef typename traits<DiagonalVectorType>::XprKind XprKind;
        enum
        {
            RowsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
            ColsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
            MaxRowsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
            MaxColsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
            Flags = (traits<DiagonalVectorType>::Flags & LvalueBit) | NoPreferredStorageOrderBit
        };
    };
}  // namespace internal

template <typename _DiagonalVectorType> class DiagonalWrapper : public DiagonalBase<DiagonalWrapper<_DiagonalVectorType>>, internal::no_assignment_operator
{
public:
#ifndef EIGEN_PARSED_BY_DOXYGEN
    typedef _DiagonalVectorType DiagonalVectorType;
    typedef DiagonalWrapper Nested;
#endif

    /** Constructor from expression of diagonal coefficients to wrap. */
    EIGEN_DEVICE_FUNC
    explicit inline DiagonalWrapper(DiagonalVectorType& a_diagonal) : m_diagonal(a_diagonal) {}

    /** \returns a const reference to the wrapped expression of diagonal coefficients. */
    EIGEN_DEVICE_FUNC
    const DiagonalVectorType& diagonal() const { return m_diagonal; }

protected:
    typename DiagonalVectorType::Nested m_diagonal;
};

/** \returns a pseudo-expression of a diagonal matrix with *this as vector of diagonal coefficients
  *
  * \only_for_vectors
  *
  * Example: \include MatrixBase_asDiagonal.cpp
  * Output: \verbinclude MatrixBase_asDiagonal.out
  *
  * \sa class DiagonalWrapper, class DiagonalMatrix, diagonal(), isDiagonal()
  **/
template <typename Derived> EIGEN_DEVICE_FUNC inline const DiagonalWrapper<const Derived> MatrixBase<Derived>::asDiagonal() const
{
    return DiagonalWrapper<const Derived>(derived());
}

/** \returns true if *this is approximately equal to a diagonal matrix,
  *          within the precision given by \a prec.
  *
  * Example: \include MatrixBase_isDiagonal.cpp
  * Output: \verbinclude MatrixBase_isDiagonal.out
  *
  * \sa asDiagonal()
  */
template <typename Derived> bool MatrixBase<Derived>::isDiagonal(const RealScalar& prec) const
{
    if (cols() != rows())
        return false;
    RealScalar maxAbsOnDiagonal = static_cast<RealScalar>(-1);
    for (Index j = 0; j < cols(); ++j)
    {
        RealScalar absOnDiagonal = numext::abs(coeff(j, j));
        if (absOnDiagonal > maxAbsOnDiagonal)
            maxAbsOnDiagonal = absOnDiagonal;
    }
    for (Index j = 0; j < cols(); ++j)
        for (Index i = 0; i < j; ++i)
        {
            if (!internal::isMuchSmallerThan(coeff(i, j), maxAbsOnDiagonal, prec))
                return false;
            if (!internal::isMuchSmallerThan(coeff(j, i), maxAbsOnDiagonal, prec))
                return false;
        }
    return true;
}

namespace internal {

    template <> struct storage_kind_to_shape<DiagonalShape>
    {
        typedef DiagonalShape Shape;
    };

    struct Diagonal2Dense
    {
    };

    template <> struct AssignmentKind<DenseShape, DiagonalShape>
    {
        typedef Diagonal2Dense Kind;
    };

    // Diagonal matrix to Dense assignment
    template <typename DstXprType, typename SrcXprType, typename Functor> struct Assignment<DstXprType, SrcXprType, Functor, Diagonal2Dense>
    {
        static void run(DstXprType& dst, const SrcXprType& src, const internal::assign_op<typename DstXprType::Scalar, typename SrcXprType::Scalar>& /*func*/)
        {
            Index dstRows = src.rows();
            Index dstCols = src.cols();
            if ((dst.rows() != dstRows) || (dst.cols() != dstCols))
                dst.resize(dstRows, dstCols);

            dst.setZero();
            dst.diagonal() = src.diagonal();
        }

        static void
        run(DstXprType& dst, const SrcXprType& src, const internal::add_assign_op<typename DstXprType::Scalar, typename SrcXprType::Scalar>& /*func*/)
        {
            dst.diagonal() += src.diagonal();
        }

        static void
        run(DstXprType& dst, const SrcXprType& src, const internal::sub_assign_op<typename DstXprType::Scalar, typename SrcXprType::Scalar>& /*func*/)
        {
            dst.diagonal() -= src.diagonal();
        }
    };

}  // namespace internal

}  // end namespace Eigen

#endif  // EIGEN_DIAGONALMATRIX_H
